# Load packages
library(glmnet)
library(magrittr)
library(selectiveInference)
library(tidyverse)

# Data
# https://office365stanford-my.sharepoint.com/:x:/r/personal/simonlev_stanford_edu/_layouts/15/Doc.aspx?sourcedoc=%7BCE7DAF19-8F1A-43B0-866D-3D0F36336EDC%7D&file=Peds%20AVM%20Database%20July%2010%202024%20Onwards.xlsx&fromShare=true&action=default&mobileredirect=true

# Source code

# Source outcome-specific configs
source(params$CONFIG)

# Destination locations
DST_DIRNAME <- paste0("../../outputs/")

# Source locations
SRC_DIRNAME <- "../../data/3_tidy/patients"
SRC_BASENAME <- "patients-daily_2024-07-21.csv"

# Should the Dockerfile be automatically updated?
UPDATE_DOCKERFILE <- FALSE

# Set the seed
set.seed(1891)

EXPOSURES_CONTINUOUS <- c(
  "Age at 1st treatment (years)" = "age_at_first_treatment_yrs",
  "mRS (presentation)" = "modified_rankin_score_presentation",
  "mRS (pre-treatment)" = "modified_rankin_score_pretreatment",
  "mRS (1-week post-op)" = "modified_rankin_score_postop_within_1_week",
  "mRS (final)" = "modified_rankin_score_final",
  "Nidus size (cm)" = "max_size_cm",
  "Spetzler-Martin grade" = "spetzler_martin_grade",
  "Lawton-Young grade" = "lawton_young_grade",
  "Size score" = "size_score"
)

EXPOSURES_BINARY <- c(
  "Sex (Male)" = "is_male",
  "Involves eloquent location" = "is_eloquent_location",
  "Has deep venous drainage" = "has_deep_venous_drainage",
  "Diffuse nidus" = "is_diffuse_nidus",
  "Hemorrhage at presentation" = "has_hemorrhage",
  "Seizures at presentation" = "has_seizures",
  "Deficit at presentation" = "has_deficit",
  "Paresis at presentation" = "has_paresis",
  "Presence of aneurysms" = "has_associated_aneurysm",
  "Spetzler-Martin grade < 4" = "is_spetzler_martin_grade_less_than_4",
  "Had surgery" = "is_surgery"
)

EXPOSURES_CATEGORICAL <- c(
  "Location" = "location",
  "Venous drainage" = "venous_drainage",
  "Modality of treatment" = "procedure_combinations"
)

OUTCOMES <- c(
  "Poor outcome (mRS >= 3)" = "is_poor_outcome",
  "Obliteration" = "is_obliterated",
  "Complications - minor" = "has_complication_minor",
  "Complications - major" = "has_complication_major",
  "mRS change (final - pre-treatment)" =
    "modified_rankin_score_final_minus_pretx",
  "mRS change (final - presentation)" =
    "modified_rankin_score_final_minus_presentation",
  "mRS change direction (Final - Pre-tx)" = "change_in_mrs_tx_vs_final"
)

plots <- list()
tables <- list()
models <- list()

patients_ <-
  patients_ %>%
  filter(!is.na(patient_id)) %>%
  arrange(patient_id)

df_uni <-
  patients_ %>%
  filter(is_eligible)

df_multi <-
  patients_ %>%
  filter(is_eligible)

df_multi <-
  df_multi %>%
  mutate(
    venous_drainage = ifelse(venous_drainage == "Both", "Deep", venous_drainage)
  )

df_multi <-
  df_multi %>%
  mutate(procedure_combinations = case_when(
    nchar(procedure_combinations) > 1 ~ "Multimodal",
    .default = procedure_combinations
  ))

df_multi <-
  df_multi %>%
  mutate(
    is_surgery = str_detect(procedure_combinations, "S")
  ) %>% 
  relocate(is_surgery, .after = procedure_combinations)

df_uni <-
  df_uni %>%
  mutate(
    is_surgery = str_detect(procedure_combinations, "S")
  ) %>% 
  relocate(is_surgery, .after = procedure_combinations)

df_multi <-
  df_multi %>%
  mutate(
    location = ifelse(location == "Corpus Callosum", "Other", location),
    location = ifelse(location == "Cerebellum", "Other", location),
    # location = ifelse(location == "Deep", "Other", location),
    location = factor(location),
    location = relevel(location, ref = "Other")
  )

# For multivariable analysis
df_multi <-
  df_multi %>%
  mutate(
    is_spetzler_martin_grade_less_than_4 = spetzler_martin_grade < 4
  ) %>% 
  relocate(is_spetzler_martin_grade_less_than_4, .after = spetzler_martin_grade)

# For univariable analysis
df_uni <-
  df_uni %>%
  mutate(
    is_spetzler_martin_grade_less_than_4 = spetzler_martin_grade < 4
  ) %>% 
  relocate(is_spetzler_martin_grade_less_than_4, .after = spetzler_martin_grade)

# Get cols
cols <- unname(c(
  EXPOSURES_CONTINUOUS,
  EXPOSURES_BINARY,
  EXPOSURES_CATEGORICAL,
  OUTCOMES
))

# Count missingness
df_multi %>%
  select(cols) %>%
  summarise_all(~ sum(is.na(.))) %>%
  pivot_longer(everything(), values_to = "num_missing") %>%
  arrange(desc(num_missing)) %>%
  filter(num_missing > 0) %>%
  sable()

df_multi %>%
  filter(is_eligible) %>%
  select(mrn, cols) %>%
  mutate(across(-mrn, is.na)) %>%
  pivot_longer(
    cols = -mrn,
    names_to = "name",
    values_to = "is_missing"
  ) %>%
  filter(is_missing) %>%
  group_by(name) %>%
  summarize(mrn = str_c(mrn, collapse = ", ")) %>%
  sable()

# Define the predictors of interest
cols <- c(
  EXPOSURES_CONTINUOUS,
  EXPOSURES_BINARY,
  CHOSEN_OUTCOME
)

# Reshape the data
df_long <-
  df_uni %>%
  select(all_of(cols)) %>%
  pivot_longer(-CHOSEN_OUTCOME, names_to = "predictor", values_to = "value")

# Plot the box plots
df_long %>%
  ggplot(aes_string(
    x = CHOSEN_OUTCOME,
    y = "value",
    fill = CHOSEN_OUTCOME,
    color = CHOSEN_OUTCOME
  )) +
  geom_boxplot(alpha = 0.5) +
  facet_wrap(~predictor, scales = "free") +
  labs(x = "Outcome", y = "Value", fill = "Outcome") +
  theme_minimal() +
  theme(
    axis.text = element_text(size = 10),
    axis.title = element_text(size = 11),
    strip.text = element_text(size = 11),
    legend.position = "none"
  )

# Define the predictors of interest
cols <- c(
  EXPOSURES_CONTINUOUS,
  EXPOSURES_BINARY,
  CHOSEN_OUTCOME
)

# Reshape the data
df_long <-
  df_uni %>%
  select(all_of(cols)) %>%
  pivot_longer(
    cols = -c(CHOSEN_OUTCOME, `Hemorrhage at presentation`),
    names_to = "predictor",
    values_to = "value"
  )

# Plot the box plots
df_long %>%
  ggplot(aes_string(
    x = CHOSEN_OUTCOME,
    y = "value",
    fill = "`Hemorrhage at presentation`",
    color = "`Hemorrhage at presentation`"
  )) +
  geom_boxplot(alpha = 0.5) +
  facet_wrap(~predictor, scales = "free") +
  labs(x = "Outcome", y = "Value", fill = "Outcome") +
  theme_minimal() +
  guides(fill = "none") +
  theme(
    axis.text = element_text(size = 10),
    axis.title = element_text(size = 11),
    strip.text = element_text(size = 11)
  )

compute_descriptive_continuous <- function(
    df = df_uni,
    cols = c(EXPOSURES_CONTINUOUS),
    use_interaction = TRUE,
    interaction_col = "has_hemorrhage",
    interaction_col_name = "Hemorrhage at presentation",
    interaction_col_value = TRUE,
    prefix = "Haem"
) {
  
  # Apply desired formatting to numbers
  format_numbers <- function(x) {
    sprintf("%.1f", x)
  }
  
  # Compute new dataset
  df <-
    df %>%
    select(cols, CHOSEN_OUTCOME, interaction_col) %>%
    drop_na(CHOSEN_OUTCOME, interaction_col)
  
  # Apply interaction
  if (use_interaction) {
    df <- 
      df %>%
      filter(!!sym(interaction_col) == interaction_col_value)
  }
  
  # Compute descriptive statistics
  desc_stats <-
    df %>%
    pivot_longer(where(is.numeric), names_to = "keys", values_to = "values") %>%
    group_by(!!sym(CHOSEN_OUTCOME), keys) %>%
    summarize(
      num_total = n(),
      num_missing = sum(is.na(values)),
      mean = mean(values, na.rm = TRUE),
      sd = sd(values, na.rm = TRUE),
      min = quantile(values, 0, na.rm = TRUE),
      max = quantile(values, 1, na.rm = TRUE),
      q_50 = quantile(values, 0.50, na.rm = TRUE),
      q_25 = quantile(values, 0.25, na.rm = TRUE),
      q_75 = quantile(values, 0.75, na.rm = TRUE)
    ) %>%
    ungroup() %>%
    mutate(
      stats = glue::glue("{format_numbers(mean)} ({format_numbers(sd)})"),
      !!sym(CHOSEN_OUTCOME) := factor(
        !!sym(CHOSEN_OUTCOME),
        levels = CHOSEN_OUTCOME_LEVELS,
        labels = CHOSEN_OUTCOME_LABELS
      )
    ) %>%
    pivot_wider(
      id_cols = keys,
      names_from = !!sym(CHOSEN_OUTCOME),
      values_from = stats
    ) %>%
    relocate(CHOSEN_OUTCOME_LABELS[2], .after = CHOSEN_OUTCOME_LABELS[1])
  
  # Compute p-values
  pvals <-
    df %>%
    pivot_longer(where(is.numeric), names_to = "keys", values_to = "values") %>%
    group_by(keys) %>%
    summarize(
      pvalue = wilcox.test(values ~ !!sym(CHOSEN_OUTCOME))$p.value
    )
  
  # Synthesize
  out <-
    desc_stats %>%
    left_join(pvals)
  
  # Add sample size to column names
  num_with_outcome <- sum(df %>% pull(CHOSEN_OUTCOME))
  num_no_outcome <- sum(!df %>% pull(CHOSEN_OUTCOME))
  
  new_col_names <- c(
    glue::glue("{prefix} - {CHOSEN_OUTCOME_LABELS[2]} (n={num_with_outcome})"),
    glue::glue("{prefix} - {CHOSEN_OUTCOME_LABELS[1]} (n={num_no_outcome})"),
    glue::glue("{prefix} - P-value")
  )
  
  out <-
    out %>%
    rename_with(
      .cols = !!sym(CHOSEN_OUTCOME_LABELS[1]):pvalue, ~ new_col_names
    ) %>%
    mutate(
      keys = ifelse(keys == interaction_col, interaction_col_name, keys)
    )
  
  return(out)
}

compute_descriptive_binary <- function(
    df = df_uni,
    cols = c(EXPOSURES_BINARY),
    use_interaction = TRUE,
    interaction_col = "has_hemorrhage",
    interaction_col_name = "Hemorrhage at presentation",
    interaction_col_value = TRUE,
    prefix = "Haem"
) {
  
  # Define arguments
  cols <- c(cols[!cols == interaction_col])
  
  # Apply desired formatting to numbers
  format_numbers <- function(x, decimals = 0) {
    decimals <- paste0("%.", decimals, "f")
    sprintf(decimals, x)
  }
  
  # Compute new dataset
  df <-
    df %>%
    select(all_of(cols), CHOSEN_OUTCOME, interaction_col) %>%
    drop_na(CHOSEN_OUTCOME, interaction_col)
  
  # Apply interaction
  if (use_interaction) {
    df <- 
      df %>%
      filter(!!sym(interaction_col) == interaction_col_value)
  }
  
  # Compute descriptive statistics
  desc_stats <-
    df %>%
    pivot_longer(
      cols = -c(CHOSEN_OUTCOME),
      names_to = "keys",
      values_to = "values"
    ) %>%
    group_by(!!sym(CHOSEN_OUTCOME), keys) %>%
    summarize(
      num_total = n(),
      num_missing = sum(is.na(values)),
      num_with = sum(values, na.rm = TRUE),
      num_without = sum(!values, na.rm = TRUE),
      pct_with = mean(values, na.rm = TRUE)
    ) %>%
    ungroup() %>%
    mutate(
      stats = glue::glue("{num_with} ({format_numbers(pct_with*100)}%)"),
      !!sym(CHOSEN_OUTCOME) := factor(
        !!sym(CHOSEN_OUTCOME),
        levels = CHOSEN_OUTCOME_LEVELS,
        labels = CHOSEN_OUTCOME_LABELS
      )
    ) %>%
    pivot_wider(
      id_cols = keys,
      names_from = !!sym(CHOSEN_OUTCOME),
      values_from = stats
    ) %>%
    relocate(CHOSEN_OUTCOME_LABELS[1], .after = CHOSEN_OUTCOME_LABELS[2])
  
  # Compute p-values
  pvals <-
    df %>%
    pivot_longer(
      cols = -c(CHOSEN_OUTCOME),
      names_to = "keys",
      values_to = "values"
    ) %>%
    group_by(keys) %>%
    summarize(
      pvalue = kruskal.test(values ~ !!sym(CHOSEN_OUTCOME))$p.value
    )
  
  # Synthesize
  out <-
    desc_stats %>%
    left_join(pvals)
  
  # Prettify
  num_with_outcome <- sum(df %>% pull(CHOSEN_OUTCOME))
  num_no_outcome <- sum(!df %>% pull(CHOSEN_OUTCOME))
  
  new_col_names <- c(
    glue::glue("{prefix} - {CHOSEN_OUTCOME_LABELS[2]} (n={num_with_outcome})"),
    glue::glue("{prefix} - {CHOSEN_OUTCOME_LABELS[1]} (n={num_no_outcome})"),
    glue::glue("{prefix} - P-value")
  )
  
  out <-
    out %>%
    rename_with(
      .cols = !!sym(CHOSEN_OUTCOME_LABELS[2]):pvalue, ~new_col_names
    ) %>%
    mutate(
      keys = ifelse(keys == interaction_col, interaction_col_name, keys)
    )
  
  return(out)
}

# Rename dataframe
`Pediatric AVMs` <-
  df_uni %>%
  filter(is_eligible) %>%
  select(-is_eligible, -comments)

# Create summary
`Pediatric AVMs` %>%
  summarytools::dfSummary(display.labels = FALSE) %>%
  print(
    file =
      "../../outputs/descriptive-statistics/descriptive_statistics_eligible.html",
    footnote = NA
  )

# Remove unwanted dataframe
remove(`Pediatric AVMs`)

# Define arguments
interaction_col = "has_hemorrhage"
interaction_col_name = "Hemorrhage at presentation"
prefix_true = "Haem"
prefix_false = "No Haem"

# Get all
out <- list()
out[[1]] <- compute_descriptive_continuous(
  use_interaction = F,
  interaction_col = interaction_col,
  interaction_col_name = interaction_col_name,
  interaction_col_value = T,
  prefix = "All"
)
out[[2]] <- compute_descriptive_binary(
  use_interaction = F,
  interaction_col = interaction_col,
  interaction_col_name = interaction_col_name,
  interaction_col_value = T,
  prefix = "All"
)
out_all <- bind_rows(out)

# Get with hemorrhage
out <- list()
out[[1]] <- compute_descriptive_continuous(
  use_interaction = T,
  interaction_col = interaction_col,
  interaction_col_name = interaction_col_name,
  interaction_col_value = T,
  prefix = prefix_true
)
out[[2]] <- compute_descriptive_binary(
  use_interaction = T,
  interaction_col = interaction_col,
  interaction_col_name = interaction_col_name,
  interaction_col_value = T,
  prefix = prefix_true
)
out_haem <- bind_rows(out)

# Get without hemorrhage
out <- list()
out[[1]] <- compute_descriptive_continuous(
  use_interaction = T,
  interaction_col = interaction_col,
  interaction_col_name = interaction_col_name,
  interaction_col_value = F,
  prefix = prefix_false
)
out[[2]] <- compute_descriptive_binary(
  use_interaction = T,
  interaction_col = interaction_col,
  interaction_col_name = interaction_col_name,
  interaction_col_value = F,
  prefix = prefix_false
)
out_no_haem <- bind_rows(out)

# Synthesize
out_complete <-
  out_all %>%
  left_join(out_haem, by = "keys") %>%
  left_join(out_no_haem, by = "keys") %>%
  arrange(`All - P-value`)

# Print
out_complete %>%
  sable()

# Define arguments
interaction_col = "is_diffuse_nidus"
interaction_col_name = "Is diffuse nidus"
prefix_true = "Diffuse"
prefix_false = "Not Diffuse"

# Get all
out <- list()
out[[1]] <- compute_descriptive_continuous(
  use_interaction = F,
  interaction_col = interaction_col,
  interaction_col_name = interaction_col_name,
  interaction_col_value = T,
  prefix = "All"
)
out[[2]] <- compute_descriptive_binary(
  use_interaction = F,
  interaction_col = interaction_col,
  interaction_col_name = interaction_col_name,
  interaction_col_value = T,
  prefix = "All"
)
out_all <- bind_rows(out)

# Get with hemorrhage
out <- list()
out[[1]] <- compute_descriptive_continuous(
  use_interaction = T,
  interaction_col = interaction_col,
  interaction_col_name = interaction_col_name,
  interaction_col_value = T,
  prefix = prefix_true
)
out[[2]] <- compute_descriptive_binary(
  use_interaction = T,
  interaction_col = interaction_col,
  interaction_col_name = interaction_col_name,
  interaction_col_value = T,
  prefix = prefix_true
)
out_haem <- bind_rows(out)

# Get without hemorrhage
out <- list()
out[[1]] <- compute_descriptive_continuous(
  use_interaction = T,
  interaction_col = interaction_col,
  interaction_col_name = interaction_col_name,
  interaction_col_value = F,
  prefix = prefix_false
)
out[[2]] <- compute_descriptive_binary(
  use_interaction = T,
  interaction_col = interaction_col,
  interaction_col_name = interaction_col_name,
  interaction_col_value = F,
  prefix = prefix_false
)
out_no_haem <- bind_rows(out)

# Synthesize
out_complete <-
  out_all %>%
  left_join(out_haem, by = "keys") %>%
  left_join(out_no_haem, by = "keys") %>%
  arrange(`All - P-value`)

# Print
out_complete %>%
  sable()

fit_model <- function(
    df = df_uni, cols = cols, is_sandwich = T) {
  # Initialize
  out <- list()

  for (i in seq_along(cols)) {
    # Create formula
    model <- as.formula(paste(CHOSEN_OUTCOME, "~", cols[i]))

    fit <- suppressWarnings(glm(
      model,
      data = df,
      family = binomial()
    ))

    if (is_sandwich) {
      # Calculate robust standard errors (sandwich)
      robust_se <- sandwich::vcovHC(fit, type = "HC0")
      # Compute robust confidence intervals
      fit_results <- lmtest::coeftest(fit, vcov. = robust_se)
    } else {
      fit_results <- fit
    }

    # Summarize model coefficients
    fit_summary <-
      fit_results %>%
      # broom does not support exponentiation after coeftest so do it manually
      broom::tidy(conf.int = T) %>%
      mutate(across(c(estimate, "conf.low", "conf.high"), exp)) %>%
      arrange(term == "(Intercept)", p.value) %>%
      rename(odds_ratio = estimate, z_value = statistic)

    # Stylize and print
    stylized <-
      fit_summary %>%
      rename(
        "Predictors" = term,
        "Odds Ratios (OR)" = odds_ratio,
        "SE" = std.error,
        "Z-scores" = z_value,
        "P-values" = p.value,
        "CI (low)" = conf.low,
        "CI (high)" = conf.high,
      ) %>%
      mutate(
        `Odds Ratios (OR)` = round(`Odds Ratios (OR)`, 2),
        `SE` = round(`SE`, 2),
        `CI (low)` = round(`CI (low)`, 2),
        `CI (high)` = round(`CI (high)`, 2),
        `P-values` = round(`P-values`, 3),
      )

    out[[i]] <- stylized
  }
  return(out)
}

# Define cols
cols <- unname(c(
  EXPOSURES_CONTINUOUS,
  EXPOSURES_BINARY,
  EXPOSURES_CATEGORICAL
))

# Fit model
out <- fit_model(df = df_uni, cols = cols, is_sandwich = T)

# Create table
univariable_unadjusted <-
  out %>%
  bind_rows() %>%
  filter(Predictors != "(Intercept)") %>%
  filter(`CI (high)` < 50) %>%
  arrange(`P-values`)

# Print
univariable_unadjusted %>%
  sable()

cols <- unname(c(
  EXPOSURES_CONTINUOUS,
  EXPOSURES_BINARY,
  EXPOSURES_CATEGORICAL
))

# Initialize
out <- list()
df <- df_uni

for (i in seq_along(cols)) {
  # Create formula
  model <- as.formula(paste(
    CHOSEN_OUTCOME,
    "~",
    "age_at_first_treatment_yrs + is_male +",
    cols[i]
  ))

  fit <- suppressWarnings(glm(
    model,
    data = df,
    family = binomial()
  ))

  # Calculate robust standard errors (sandwich)
  robust_se <- sandwich::vcovHC(fit, type = "HC0")
  # Compute robust confidence intervals
  robust_ci <- lmtest::coeftest(fit, vcov. = robust_se)

  # Summarize model coefficients
  fit_summary <-
    robust_ci %>%
    # broom does not support exponentiation after coeftest so do it manually
    broom::tidy(conf.int = T) %>%
    mutate(across(c(estimate, "conf.low", "conf.high"), exp)) %>%
    arrange(term == "(Intercept)", p.value) %>%
    rename(odds_ratio = estimate, z_value = statistic)

  # Stylize and print
  stylized <-
    fit_summary %>%
    rename(
      "Predictors" = term,
      "Odds Ratios (OR)" = odds_ratio,
      "SE" = std.error,
      "Z-scores" = z_value,
      "P-values" = p.value,
      "CI (low)" = conf.low,
      "CI (high)" = conf.high,
    ) %>%
    mutate(
      `Odds Ratios (OR)` = round(`Odds Ratios (OR)`, 2),
      `SE` = round(`SE`, 2),
      `CI (low)` = round(`CI (low)`, 2),
      `CI (high)` = round(`CI (high)`, 2),
      `P-values` = round(`P-values`, 3),
    )

  out[[i]] <- stylized
}

# Create table
univariable_adjusted <-
  out %>%
  bind_rows() %>%
  filter(Predictors != "(Intercept)") %>%
  filter(!str_starts(Predictors, "age_at_first_treatment_yrs|is_male")) %>%
  filter(`CI (high)` < 50) %>%
  arrange(`P-values`)

# Print
univariable_adjusted %>%
  sable()

# Define columns
cols <- unname(c(
  EXPOSURES_CONTINUOUS,
  EXPOSURES_BINARY,
  EXPOSURES_CATEGORICAL
))

# Initialize
out <- list()
df <- df_multi

for (i in seq_along(cols)) {
  # Create formula
  model <- as.formula(paste(CHOSEN_OUTCOME, "~", cols[i]))

  fit <- suppressWarnings(glm(
    model,
    data = df,
    family = binomial()
  ))

  # Calculate robust standard errors (sandwich)
  robust_se <- sandwich::vcovHC(fit, type = "HC0")
  # Compute robust confidence intervals
  robust_ci <- lmtest::coeftest(fit, vcov. = robust_se)

  # Summarize model coefficients
  fit_summary <-
    robust_ci %>%
    # broom does not support exponentiation after coeftest so do it manually
    broom::tidy(conf.int = T) %>%
    mutate(across(c(estimate, "conf.low", "conf.high"), exp)) %>%
    arrange(term == "(Intercept)", p.value) %>%
    rename(odds_ratio = estimate, z_value = statistic)

  # Stylize and print
  stylized <-
    fit_summary %>%
    rename(
      "Predictors" = term,
      "Odds Ratios (OR)" = odds_ratio,
      "SE" = std.error,
      "Z-scores" = z_value,
      "P-values" = p.value,
      "CI (low)" = conf.low,
      "CI (high)" = conf.high,
    ) %>%
    mutate(
      `Odds Ratios (OR)` = round(`Odds Ratios (OR)`, 2),
      `SE` = round(`SE`, 2),
      `CI (low)` = round(`CI (low)`, 2),
      `CI (high)` = round(`CI (high)`, 2),
      `P-values` = round(`P-values`, 3),
    )

  out[[i]] <- stylized
}

# Create table
univariable_unadjusted_recoded <-
  out %>%
  bind_rows() %>%
  filter(Predictors != "(Intercept)") %>%
  filter(`CI (high)` < 50) %>%
  arrange(`P-values`)

# Print
univariable_unadjusted_recoded %>%
  sable()

# Define columns
cols <- unname(c(
  EXPOSURES_CONTINUOUS,
  EXPOSURES_BINARY,
  EXPOSURES_CATEGORICAL
))

# Initialize
out <- list()
df <- df_multi

for (i in seq_along(cols)) {
  # Create formula
  model <- as.formula(paste(
    CHOSEN_OUTCOME,
    "~",
    "age_at_first_treatment_yrs + is_male +",
    cols[i]
  ))

  fit <- suppressWarnings(glm(
    model,
    data = df,
    family = binomial()
  ))

  # Calculate robust standard errors (sandwich)
  robust_se <- sandwich::vcovHC(fit, type = "HC0")
  # Compute robust confidence intervals
  robust_ci <- lmtest::coeftest(fit, vcov. = robust_se)

  # Summarize model coefficients
  fit_summary <-
    robust_ci %>%
    # broom does not support exponentiation after coeftest so do it manually
    broom::tidy(conf.int = T) %>%
    mutate(across(c(estimate, "conf.low", "conf.high"), exp)) %>%
    arrange(term == "(Intercept)", p.value) %>%
    rename(odds_ratio = estimate, z_value = statistic)

  # Stylize and print
  stylized <-
    fit_summary %>%
    rename(
      "Predictors" = term,
      "Odds Ratios (OR)" = odds_ratio,
      "SE" = std.error,
      "Z-scores" = z_value,
      "P-values" = p.value,
      "CI (low)" = conf.low,
      "CI (high)" = conf.high,
    ) %>%
    mutate(
      `Odds Ratios (OR)` = round(`Odds Ratios (OR)`, 2),
      `SE` = round(`SE`, 2),
      `CI (low)` = round(`CI (low)`, 2),
      `CI (high)` = round(`CI (high)`, 2),
      `P-values` = round(`P-values`, 3),
    )

  out[[i]] <- stylized
}

# Create table
univariable_adjusted_recoded <-
  out %>%
  bind_rows() %>%
  filter(Predictors != "(Intercept)") %>%
  filter(!str_starts(Predictors, "age_at_first_treatment_yrs|is_male")) %>%
  filter(`CI (high)` < 50) %>%
  arrange(`P-values`)

# Print
univariable_adjusted_recoded %>%
  sable()

# Define columns
cols <- unname(c(
  EXPOSURES_CONTINUOUS,
  EXPOSURES_BINARY,
  EXPOSURES_CATEGORICAL
))

# Remove unwanted columns
unwanted <- c("modified_rankin_score_final")
cols <- cols[!cols %in% unwanted]

adjustors <- c(
  "age_at_first_treatment_yrs",
  "is_male"
)

# Initialize
out <- list()
df <- df_multi
k <- 1

for (i in seq_along(cols)) {
  # Do not fit model for variables that we are adjusting for
  if (cols[i] %in% adjustors) {
    next
  }

  # Create formula
  model <- as.formula(paste(
    CHOSEN_OUTCOME,
    "~",
    "age_at_first_treatment_yrs + is_male +",
    cols[i],
    "* has_hemorrhage",
    sep = ""
  ))

  fit <- suppressWarnings(glm(
    model,
    data = df,
    family = binomial()
  ))

  # Calculate robust standard errors (sandwich)
  robust_se <- sandwich::vcovHC(fit, type = "HC0")
  # Compute robust confidence intervals
  robust_ci <- lmtest::coeftest(fit, vcov. = robust_se)

  # Summarize model coefficients
  fit_summary <-
    robust_ci %>%
    # broom does not support exponentiation after coeftest so do it manually
    broom::tidy(conf.int = T) %>%
    mutate(across(c(estimate, "conf.low", "conf.high"), exp)) %>%
    arrange(term == "(Intercept)", p.value) %>%
    rename(odds_ratio = estimate, z_value = statistic)

  # Stylize and print
  stylized <-
    fit_summary %>%
    rename(
      "Predictors" = term,
      "Odds Ratios (OR)" = odds_ratio,
      "SE" = std.error,
      "Z-scores" = z_value,
      "P-values" = p.value,
      "CI (low)" = conf.low,
      "CI (high)" = conf.high,
    ) %>%
    mutate(
      `Odds Ratios (OR)` = round(`Odds Ratios (OR)`, 2),
      `SE` = round(`SE`, 2),
      `CI (low)` = round(`CI (low)`, 2),
      `CI (high)` = round(`CI (high)`, 2),
      `P-values` = round(`P-values`, 3),
    )

  out[[k]] <- stylized
  k <- k + 1
}

# Define values
unwanted <- c(
  "is_maleTRUE",
  "age_at_first_treatment_yrs",
  "(Intercept)",
  "has_hemorrhageTRUE"
)

# Initialize objects
isolated_values <- list()

# Get main and interaction effects
for (i in seq_along(out)) {
  # Get data
  result <- out[[i]]

  # Get all main effects of interest
  main_effects <-
    result %>%
    filter(!Predictors %in% unwanted) %>%
    filter(!str_detect(Predictors, ":")) %>%
    select(-"SE", -"Z-scores")

  # Get the interaction effects
  interaction_effects <-
    result %>%
    filter(!Predictors %in% unwanted) %>%
    filter(str_detect(Predictors, ":")) %>%
    mutate(Predictors = str_extract(Predictors, "^[^:]+")) %>%
    select(-"SE", -"Z-scores") %>%
    rename_with(~ paste("Interaction -", .), -Predictors)

  isolated_values[[i]] <-
    main_effects %>%
    left_join(interaction_effects, by = "Predictors")
}

# Create table
univariable_adjusted_recoded_interactions <-
  isolated_values %>%
  bind_rows() %>%
  arrange(`Interaction - P-values`)

# Print
univariable_adjusted_recoded_interactions %>%
  sable()

df %>%
  count(has_hemorrhage, is_eloquent_location, !!sym(CHOSEN_OUTCOME)) %>%
  drop_na() %>%
  arrange(has_hemorrhage, is_eloquent_location) %>%
  group_by(has_hemorrhage, is_eloquent_location) %>%
  mutate(
    num_total = sum(n),
    pct_total = scales::percent(n / num_total, 1)
  ) %>%
  sable()

df %>%
  count(has_hemorrhage, has_deficit, !!sym(CHOSEN_OUTCOME)) %>%
  drop_na() %>%
  arrange(has_hemorrhage, has_deficit) %>%
  group_by(has_hemorrhage, has_deficit) %>%
  mutate(
    num_total = sum(n),
    pct_total = scales::percent(n / num_total, 1)
  ) %>%
  sable()

df %>%
  count(has_hemorrhage, has_seizures, !!sym(CHOSEN_OUTCOME)) %>%
  drop_na() %>%
  arrange(has_hemorrhage, has_seizures) %>%
  group_by(has_hemorrhage, has_seizures) %>%
  mutate(
    num_total = sum(n),
    pct_total = scales::percent(n / num_total, 1)
  ) %>%
  sable()

unwanted_all <- c(
  # Use values at or before first treatment
  "modified_rankin_score_postop_within_1_week",
  "modified_rankin_score_final"
)
  
unwanted_without_gradings <- c(
  unwanted_all,
  # Spetzler-Martin grade (size score + eloquent location + venous drainage)
  "spetzler_martin_grade",
  "is_spetzler_martin_grade_less_than_4",
  "size_score", # Already covered by the more detailed max_size_cm
  "venous_drainage", # Already covered by has_deep_venous_drainage
  "location", # Already covered to some extent by eloquence
  # Lawton-Young grade (SM + age + diffuse nidus + hemorrhage)
  # (SM is heavily overlapping with LY, so include its components)
  "lawton_young_grade",
  # Use mRS pre-treatment (more predictive)
  "modified_rankin_score_presentation"  # "modified_rankin_score_pretreatment"
)

# Define cols of interest - use grading scores whenever these exist
unwanted_with_gradings <- c(
  # Spetzler-Martin grade (size score + eloquent location + venous drainage)
  "is_spetzler_martin_grade_less_than_4",  # The grade has lower p-value
  "size_score",
  "max_size_cm",
  "is_eloquent_location",
  # "location",  # Include this as not highly correlated with SM
  "has_deep_venous_drainage",
  "venous_drainage",
  # Lawton-Young grade (SM + age + diffuse nidus + hemorrhage)
  # (SM is heavily overlapping with LY, so include its components)
  "lawton_young_grade",
  # "is_diffuse_nidus",
  # "has_hemorrhage",
  # Use values at or before first treatment
  "modified_rankin_score_final",
  "modified_rankin_score_postop_within_1_week",
  # Use mRS pre-treatment (more predictive)
  "modified_rankin_score_presentation"
  # "modified_rankin_score_pretreatment"
)

# Define cols
cols <- unname(c(
  EXPOSURES_CONTINUOUS,
  EXPOSURES_BINARY,
  EXPOSURES_CATEGORICAL,
  CHOSEN_OUTCOME
))

# Define column sets for each analysis
cols_all <- cols[!cols %in% unwanted_all]
cols_with_gradings <- cols[!cols %in% unwanted_with_gradings]
cols_without_gradings <- cols[!cols %in% unwanted_without_gradings]

# Create df of interest
df_all <-
  # Use the non-recoded dataset
  df_uni %>%
  dplyr::select(all_of(cols_all)) %>%
  drop_na()

df_with_gradings <-
  # Use the non-recoded dataset
  df_uni %>%
  dplyr::select(all_of(cols_with_gradings)) %>%
  drop_na()

df_no_scores <-
  # Use the non-recoded dataset
  df_uni %>%
  dplyr::select(all_of(cols_without_gradings)) %>%
  drop_na()

# Create formula
frla <- as.formula(paste(CHOSEN_OUTCOME, " ~ . - 1"))

# Create matrices with all variables
X_all <- model.matrix(frla, df_all)
y_all <- df_all %>% pull(CHOSEN_OUTCOME) %>% as.numeric()

# Create matrices with scores
X_with_gradings <- model.matrix(frla, df_with_gradings)
y_with_gradings <- df_with_gradings %>% pull(CHOSEN_OUTCOME) %>% as.numeric()

# Create matrices with score components
X_without_gradings <- model.matrix(frla, df_no_scores)
y_without_gradings <- df_no_scores %>% pull(CHOSEN_OUTCOME) %>% as.numeric()

# X should be centered for LASSO (necessary for glmnet)
# (scale = FALSE as this is done by the standardization step in glmnet)
X_all_scaled <- scale(X_all, center = T, scale = F)
X_with_gradings_scaled <- scale(X_with_gradings, center = T, scale = F)
X_without_gradings_scaled <- scale(X_without_gradings, center = T, scale = F)

# See the names of X to match column numbers to column names later on
colnames(X_without_gradings_scaled)

# Set seed
set.seed(33)

# Run forward step-wise
fsfit <- fs(
  x = X_without_gradings_scaled,
  y = y_without_gradings,
  maxsteps = 2000,
  intercept = TRUE,
  normalize = FALSE
)

# Estimate sigma
sigmahat <- estimateSigma(
  x = X_without_gradings_scaled,
  y = y_without_gradings,
  intercept = TRUE,
  standardize = FALSE
)$sigmahat

# Compute sequential p-values and confidence intervals - for all
# (sigma estimated from full model)
out.seq <- fsInf(fsfit, type = "active", sigma = sigmahat)
out.seq

# Compute p-values and confidence intervals after AIC stopping - for selected
out.aic <- fsInf(fsfit, type = "aic", sigma = sigmahat)
out.aic

# Compute p-values and confidence intervals after 5 fixed steps
out.fix <- fsInf(fsfit, type = "all", k = 5, sigma = sigmahat)
out.fix

# Set seed
set.seed(141845)

# Define values
X <- X_all
y <- y_all

# Perform cross-validation
# (alpha = 1 for lasso, alpha = 0 for ridge, 0 < alpha < 1 for elastic net)
cv_fit <- cv.glmnet(
  x = X,
  y = y,
  family = "binomial",
  alpha = 1,
  nfolds = 10
)

# Extract the best lambda
best_lambda_min <- cv_fit$lambda.min
best_lambda_1se <- cv_fit$lambda.1se

# Run glmnet
gfit <- glmnet(
  x = X,
  y = y,
  family = "binomial",
  alpha = 1,
  intercept = TRUE,
  standardize = TRUE
)

# Extract coef for a given lambda - avoid the intercept term
# (given the small number of observations hence large SD, use min lambda)
lambda <- best_lambda_min
n <- nrow(y)
beta <- coef(
  x = X,
  y = y,
  object = gfit,
  s = lambda,
  exact = TRUE
)

# Estimate sigma
gsigmahat <- estimateSigma(
  x = X,
  y = y,
  intercept = TRUE,
  standardize = TRUE
)$sigmahat

# Compute fixed lambda p-values and selection intervals
out <- fixedLassoInf(
  x = X,
  y = y,
  beta = beta,
  lambda = lambda,
  # Level of significance
  alpha = 0.1,
  family = "binomial",
  intercept = TRUE,
  sigma = gsigmahat
)

# Prettify
tibble(
  names = names(out$vars),
  odds_ratio = exp(out$coef0),
  pvalue = out$pv,
  ci_lo = exp(out$ci[, 1]),
  ci_hi = exp(out$ci[, 2])
) %>%
  arrange(pvalue) %>%
  sable()

# Set seed
set.seed(141845)

# Define values
X <- X_without_gradings
y <- y_without_gradings

# Perform cross-validation
# (alpha = 1 for lasso, alpha = 0 for ridge, 0 < alpha < 1 for elastic net)
cv_fit <- cv.glmnet(
  x = X,
  y = y,
  family = "binomial",
  alpha = 1,
  nfolds = 10
)

# Extract the best lambda
best_lambda_min <- cv_fit$lambda.min
best_lambda_1se <- cv_fit$lambda.1se

# Run glmnet
gfit <- glmnet(
  x = X,
  y = y,
  family = "binomial",
  alpha = 1,
  intercept = TRUE,
  standardize = TRUE
)

# Extract coef for a given lambda - avoid the intercept term
# (given the small number of observations hence large SD, use min lambda)
lambda <- best_lambda_min
n <- nrow(y)
beta <- coef(
  x = X,
  y = y,
  object = gfit,
  s = lambda,
  exact = TRUE
)

# Estimate sigma
gsigmahat <- estimateSigma(
  x = X,
  y = y,
  intercept = TRUE,
  standardize = TRUE
)$sigmahat

# Compute fixed lambda p-values and selection intervals
out <- fixedLassoInf(
  x = X,
  y = y,
  beta = beta,
  lambda = lambda,
  # Level of significance
  alpha = 0.1,
  family = "binomial",
  intercept = TRUE,
  sigma = gsigmahat
)

# Prettify
tibble(
  names = names(out$vars),
  odds_ratio = exp(out$coef0),
  pvalue = out$pv,
  ci_lo = exp(out$ci[, 1]),
  ci_hi = exp(out$ci[, 2])
) %>%
  arrange(pvalue) %>%
  sable()